# 落表tmp_customer_rollback_balance_special：单独计算福建南美佳美进出口贸易有限公司垫资回滚（非集团客户维度，但后期可额外做汇总）
# 三种分类：全量、优合、非金融仓
import datetime
import pandas as pd
import sys
sys.path.append("../")
sys.path.append("./")
from managers.database_manager import db_connector
from datetime import timedelta
from datetime import datetime
from sqlalchemy import text
import warnings
warnings.filterwarnings('ignore')

class DataConfig:

    FIELD_MAPPING = {
        'data_classify': '数据分类',
        'cust_code': '客户编码',
        'cust_name': '客户名称',
        'rollback_date': '回滚日期',
        'current_principal_balance': '当前本金余额',
        'month_repayment_amount': '当月还款金额',
        'month_loan_amount': '当月放款金额',
        'inv_company_code': '集团客户编码',
        'inv_company_name': '集团客户名称'
    }

    # @staticmethod
    # def get_loan_code(engine):
    #     base_sql = f"""
    #         select
    #           *
    #         FROM (
    #           SELECT
    #             cust_code,
    #             cust_name
    #           FROM dwh_dwd.dwd_order_loan_detail
    #           where is_deleted = 0
    #             and (cust_code is null or cust_name is null)
    #           GROUP BY 1,2
    #         ) t1
    #         LEFT JOIN (
    #           SELECT
    #             cust_code as cust_code_selected,
    #             cust_name as cust_name_selected
    #           FROM dwh_dwd.dwd_order_loan_detail
    #           where cust_code is not null
    #           GROUP BY 1,2
    #         ) t2
    #         on t1.cust_name = t2.cust_name_selected
    #     """
    #     df = pd.read_sql_query(base_sql, engine)
    #     return df[['cust_code_selected', 'cust_name_selected']]

    # @staticmethod
    # def get_repay_code(engine):
    #     base_sql = f"""
    #         select
    #           *
    #         FROM (
    #           SELECT
    #             cust_code,
    #             cust_name
    #           FROM dwh_dwd.dwd_repayment_flow
    #           where is_deleted = 0
    #             and (cust_code is null or cust_name is null)
    #             and cust_name not in ('（亿万银河）金百顺食品贸易（上海）有限公司', '郑州铭泰进出口有限公司') -- 剔除后全都对上了
    #           GROUP BY 1,2
    #         ) t1
    #         LEFT JOIN (
    #           SELECT
    #             cust_code as cust_code_selected,
    #             cust_name as cust_name_selected
    #           FROM dwh_dwd.dwd_order_loan_detail
    #           where cust_code is not null
    #           GROUP BY 1,2
    #         ) t2
    #         on t1.cust_name = t2.cust_name_selected
    #     """
    #     df = pd.read_sql_query(base_sql, engine)
    #     return df[['cust_code_selected', 'cust_name_selected']]

    @staticmethod
    def get_company_code(engine):
        base_sql = f"""
            SELECT
                cust_code,
                cust_name,
                IF(inv_company_code IS NULL, cust_code, inv_company_code) AS inv_company_code,
                IF(inv_company_name IS NULL, cust_name, inv_company_name) AS inv_company_name
            FROM dwh_dwd.dwh_dwd_uchain_customer_info
        """
        df = pd.read_sql_query(base_sql,engine)
        return df

# 数据查询类
class DataQuery:
    """数据查询类"""

    def __init__(self, conn_dwd, start_date=None, end_date=None):
        self.engine = conn_dwd
        self.start_date = start_date
        self.end_date = end_date

    def query_current_balance(self, finance_flag=0, loan_sign_filter=None):
        # %% 获取当前时点的客户余额（基准数据，关联集团）
        # 是否分融资类型
        subquery_finance = "-- t1.finance_type_desc,"
        group_finance = ";    -- finance_type_desc"
        if finance_flag:
            subquery_finance = "t1.finance_type_desc,"
            group_finance = ", 3;"
        # 收集子查询的条件列表
        subquery_conditions = []
        if loan_sign_filter == '优合':
            subquery_conditions.append(f"t1.loan_sign IN ('优合', '小贷', '小代采')")
        elif loan_sign_filter == '非金融仓':
            subquery_conditions.append(f"t1.loan_sign NOT IN ('金融仓')")
        conditions_str = " AND ".join(subquery_conditions)
        subquery_conditions_str = f"AND {conditions_str}" if conditions_str else ""

        # base_sql = f"""
        #     SELECT
        #       if(t1.cust_code is null, t2.cust_code, t1.cust_code) as cust_code,
        #       if(t1.cust_name is null, t2.cust_name, t1.cust_name) as cust_name,
        #       {subquery_finance}
        #       sum(t1.principal_balance) as current_principal_balance
        #     FROM dwh_dwd.dwd_order_loan_detail t1
        #     LEFT JOIN (
        #        SELECT
        #           cust_code as cust_code,
        #           cust_name as cust_name
        #        FROM dwh_dwd.dwd_order_loan_detail
        #        where cust_code is not null
        #        GROUP BY 1,2
        #     ) t2 on t1.cust_name = t2.cust_name
        #     WHERE t1.is_deleted = 0
        #       and t1.channel_flag = 0
        #       and (t1.finance_type_desc not in ('国内服务费') or t1.finance_type_desc is null)
        #       {subquery_conditions_str}
        #     GROUP BY 1,2{group_finance}
        #     """
        base_sql = f"""
            WITH
            
            -- 1.获取南美的融资单
            ext_nm_loan_dwd as (
            SELECT
              loan.*,
              analysis.cust_code as cust_code1,
              analysis.spl_code as spl_code,
              analysis.spl_name as spl_name
            -- loan 表
            FROM (
              SELECT
                t1.loan_code,
                t1.receipt_code,
                t1.order_id,
                t1.loan_start_date,
                t1.appr_loan_amt,
                t1.principal_balance,
                if(t1.cust_code is null, t2.cust_code, t1.cust_code) as cust_code,
                if(t1.cust_name is null, t2.cust_name, t1.cust_name) as cust_name
              FROM dwh_dwd.dwd_order_loan_detail t1
              LEFT JOIN (
                 SELECT
                    cust_code as cust_code,
                    cust_name as cust_name
                 FROM dwh_dwd.dwd_order_loan_detail
                 where cust_code is not null
                 GROUP BY 1,2
              ) t2 on t1.cust_name = t2.cust_name
              WHERE t1.is_deleted = 0
                and t1.channel_flag = 0
                and (t1.finance_type_desc not in ('国内服务费') or t1.finance_type_desc is null)
                {subquery_conditions_str}
            ) loan
            LEFT JOIN (
              SELECT
                order_id,
                cust_code,
                spl_code,
                spl_name
              FROM dwh_ads.ads_erp_sap_order_analysis
              where cust_code = '0020006969'
            ) analysis
              on loan.order_id = analysis.order_id
            WHERE loan.cust_code = '0020006969'
            )
            SELECT
              cust_code,
              cust_name,
              spl_code,
              spl_name,
              sum(principal_balance) as current_principal_balance
            FROM ext_nm_loan_dwd
            GROUP BY 1,2,3,4
        """

        # 替换SQL中的占位符，生成最终SQL
        final_sql = base_sql.format(
            subquery_finance=subquery_finance,
            subquery_conditions_str=subquery_conditions_str,
            group_finance=group_finance
        )

        return pd.read_sql_query(final_sql, self.engine)

    def get_month_transactions(self, finance_flag=0, loan_sign_filter=None):

        # 是否分融资类型
        subquery_finance = "-- t1.finance_type_desc,"
        group_finance = ";    -- finance_type_desc"
        if finance_flag:
            subquery_finance = "t1.finance_type_desc,"
            group_finance = ", 3;"
        # 收集子查询的条件列表
        subquery_conditions_loan = []
        subquery_conditions_repay = []
        if loan_sign_filter == '优合':
            subquery_conditions_loan.append(f"t1.loan_sign IN ('优合', '小贷', '小代采')")
            subquery_conditions_repay.append(f"t1.loan_sign IN ('优合', '小贷', '小代采')")
        elif loan_sign_filter == '非金融仓':
            subquery_conditions_loan.append(f"t1.loan_sign NOT IN ('金融仓')")
            subquery_conditions_repay.append(f"(t1.loan_sign NOT IN ('金融仓') or t1.loan_sign is null)")
        conditions_loan_str = " AND ".join(subquery_conditions_loan)
        subquery_conditions_loan_str = f"AND {conditions_loan_str}" if conditions_loan_str else ""
        conditions_repay_str = " AND ".join(subquery_conditions_repay)
        subquery_conditions_repay_str = f"AND {conditions_repay_str}" if conditions_repay_str else ""


        """获取指定月份的交易流水"""

        # 期间放款
        loan_sql = f"""
            WITH
            ext_nm_appr_dwd as (
            SELECT
              loan.*,
              analysis.cust_code as cust_code1,
              analysis.spl_code as spl_code,
              analysis.spl_name as spl_name
            -- loan 表
            FROM (
              SELECT
                t1.loan_code,
                t1.receipt_code,
                t1.order_id,
                t1.loan_start_date,
                t1.appr_loan_amt,
                t1.principal_balance,
                if(t1.cust_code is null, t2.cust_code, t1.cust_code) as cust_code,
                if(t1.cust_name is null, t2.cust_name, t1.cust_name) as cust_name
              FROM dwh_dwd.dwd_order_loan_detail t1
              LEFT JOIN (
                 SELECT
                    cust_code as cust_code,
                    cust_name as cust_name
                 FROM dwh_dwd.dwd_order_loan_detail
                 where cust_code is not null
                 GROUP BY 1,2
              ) t2 on t1.cust_name = t2.cust_name
              WHERE t1.is_deleted = 0
                and t1.channel_flag = 0
                and (t1.finance_type_desc not in ('国内服务费') or t1.finance_type_desc is null)
                AND t1.loan_start_date >= '{self.start_date.strftime('%Y-%m-%d')}'
                AND t1.loan_start_date < '{self.end_date.strftime('%Y-%m-%d')}'    
                -- AND t1.loan_start_date >= '2025-10-01'
                -- AND t1.loan_start_date < '2025-10-25'
                {subquery_conditions_loan_str}
            ) loan
            LEFT JOIN (
              SELECT
                order_id,
                cust_code,
                spl_code,
                spl_name
              FROM dwh_ads.ads_erp_sap_order_analysis
              where cust_code = '0020006969'
            ) analysis
              on loan.order_id = analysis.order_id
            WHERE loan.cust_code = '0020006969'
            )
            SELECT
              cust_code,
              cust_name,
              spl_code,
              spl_name,
              sum(appr_loan_amt) as month_loan_amount
            FROM ext_nm_appr_dwd
            GROUP BY 1,2,3,4
        """

        # 期间还款
        repayment_sql = f"""
            with 
            
            ext_nm_repayment_dwd as (
              SELECT
                loan.*,
                analysis.cust_code as cust_code1,
                analysis.spl_code as spl_code,
                analysis.spl_name as spl_name
              FROM (
                select
                  if(t1.cust_code is null, t2.cust_code, t1.cust_code) as cust_code,
                  if(t1.cust_name is null, t2.cust_name, t1.cust_name) as cust_name,
                  t1.order_id,
                  t1.loan_start_date,
                  if(t1.repayment_date < t1.loan_start_date, t1.loan_start_date, t1.repayment_date) as repayment_date,
                  t1.repayment_principal,
                  t1.finance_type,
                  t1.finance_type_desc,
                  t1.loan_sign,
                  t1.is_deleted
                FROM dwh_dwd.dwd_repayment_flow t1
                LEFT JOIN (
                   SELECT
                      cust_code as cust_code,
                      cust_name as cust_name
                   FROM dwh_dwd.dwd_order_loan_detail
                   where cust_code is not null
                   GROUP BY 1,2
                ) t2 on t1.cust_name = t2.cust_name
                where t1.is_deleted = 0
                    and t1.loan_code in (
                      SELECT loan_code FROM dwh_dwd.dwd_order_loan_detail where is_deleted = 0 and channel_flag = 0
                    )
                    AND (t1.finance_type_desc not in ('国内服务费') or t1.finance_type_desc is null)
                    and t1.cust_name not in ('（亿万银河）金百顺食品贸易（上海）有限公司', '郑州铭泰进出口有限公司')    -- 无融资单，啥情况
                    AND t1.repayment_date >= '{self.start_date.strftime('%Y-%m-%d')}'
                    AND t1.repayment_date < '{self.end_date.strftime('%Y-%m-%d')}'
                    {subquery_conditions_repay_str}
                    -- AND t1.repayment_date >= '2025-10-01'
                    -- AND t1.repayment_date < '2025-10-25'
               ) loan
               LEFT JOIN (
                 SELECT
                    order_id,
                    cust_code,
                    spl_code,
                    spl_name
                  FROM dwh_ads.ads_erp_sap_order_analysis
                  where cust_code = '0020006969'
               ) analysis
               on loan.order_id = analysis.order_id
               WHERE loan.cust_code = '0020006969'
            )
            SELECT
              cust_code,
              cust_name,
              spl_code,
              spl_name,
              sum(repayment_principal) as month_repayment_amount
            FROM ext_nm_repayment_dwd
            GROUP BY 1,2,3,4
        """

        loans = pd.read_sql_query(loan_sql, self.engine)
        repayments = pd.read_sql_query(repayment_sql, self.engine)

        return loans, repayments

def merge_with_name_fallback(monthly_balance, month_transactions_loan, month_transactions_repayment):
    # --------------------------
    # 第一步：先按 code 精确匹配（cust_code + inv_company_code）
    # --------------------------
    # 1.1 月度余额 ↔ 贷款交易（按 code 匹配）
    merged_code_loan = pd.merge(
        monthly_balance,
        month_transactions_loan,
        on=['cust_code', 'inv_company_code'],  # 仅用 code 关联（name 可能重复，暂不用于精确匹配）
        how='left',
        suffixes=('', '_loan')  # 避免列名冲突（若有同名字段）
    )

    # 1.2 再关联还款交易（按 code 匹配）
    merged_code_all = pd.merge(
        merged_code_loan,
        month_transactions_repayment,
        on=['cust_code', 'inv_company_code'],
        how='left',
        suffixes=('', '_repay')
    )

    # 标记“按 code 未匹配到”的记录（贷款或还款任一为空，即视为未匹配）
    code_unmatched_mask = merged_code_all['repayment_principal'].isna() | merged_code_all['loan_principal'].isna()
    code_matched = merged_code_all[~code_unmatched_mask].copy()  # 按 code 匹配成功的部分
    code_unmatched = merged_code_all[code_unmatched_mask].copy()  # 需用 name 兜底的部分


    # --------------------------
    # 第二步：对 code 未匹配的记录，用 name 匹配（cust_name + inv_company_name）
    # --------------------------
    if not code_unmatched.empty:
        # 2.1 提取未匹配记录的“name 键”，并清理重复（避免一对多匹配导致数据膨胀）
        name_keys = code_unmatched[['cust_name', 'inv_company_name']].drop_duplicates()

        # 2.2 name 匹配贷款交易
        merged_name_loan = pd.merge(
            name_keys,
            month_transactions_loan,
            on=['cust_name', 'inv_company_name'],  # 用 name 关联
            how='left'
        )

        # 2.3 再匹配还款交易
        merged_name_all = pd.merge(
            merged_name_loan,
            month_transactions_repayment,
            on=['cust_name', 'inv_company_name'],
            how='left'
        )

        # 2.4 将 name 匹配的结果，补充到 code 未匹配的记录中（按 name 对齐）
        code_unmatched = pd.merge(
            code_unmatched,
            merged_name_all,
            on=['cust_name', 'inv_company_name'],
            how='left',
            suffixes=('', '_name')
        )

        # 2.5 合并同名字段（用 name 匹配的结果覆盖 code 匹配的空值）
        for col in ['loan_principal', 'repayment_principal']:  # 需合并的数值字段
            code_unmatched[col] = code_unmatched[col].combine_first(code_unmatched[f'{col}_name'])
            code_unmatched.drop(columns=[f'{col}_name'], inplace=True)  # 删除临时列


    # --------------------------
    # 第三步：合并所有结果（code 匹配 + name 匹配）
    # --------------------------
    final_merged = pd.concat([code_matched, code_unmatched], ignore_index=True)

    # 填充剩余空值（数值型字段用 0，字符型字段用空字符串，避免后续计算错误）
    numeric_cols = final_merged.select_dtypes(include=['float64', 'int64']).columns
    final_merged[numeric_cols] = final_merged[numeric_cols].fillna(0)
    string_cols = final_merged.select_dtypes(include=['object']).columns
    final_merged[string_cols] = final_merged[string_cols].fillna('')

    return final_merged

def calculate_rollback(current_balance, current_date, rollback_months, conn_dwd, f, l):
    if f:
        print("\n此次回滚分融资类型 —— ——\n")
    else:
        print("\n此次回滚不分融资类型 —— ——\n")

    # 初始化结果
    results = []
    monthly_balance = current_balance.copy()


    # 从当前月份开始，倒推计算每个月月底的余额
    for i in range(len(rollback_months)):
        if i == 0:
            # 第一个月：当前实际日期
            target_date = rollback_months[-1]  # 当前实际日期
            month_start = datetime(target_date.year, target_date.month, 1)
            next_month_start = current_date + timedelta(days=1)
        else:
            # 后续月份：倒推
            target_date = rollback_months[-(i + 1)]  # 上一个月月底
            month_start = datetime(target_date.year, target_date.month, 1)
            next_month_start = datetime(target_date.year, target_date.month + 1,
                                        1) if target_date.month < 12 else datetime(target_date.year + 1, 1, 1)

        print(f"\n计算 {target_date.strftime('%Y-%m-%d')} 的回滚余额...")
        print("start_date: ", month_start)
        print("end_date: ", next_month_start)



        # 获取该月的交易数据（按月独立计算）
        query = DataQuery(conn_dwd, start_date=month_start, end_date=next_month_start)
        month_transactions_loan, month_transactions_repayment = query.get_month_transactions(finance_flag=f, loan_sign_filter=l)


        # 计算回滚余额
        if not monthly_balance.empty:
            # 合并当前余额和交易数据（使用left确保每个客户都能看到）
            # 第一步：将monthly_balance与month_transactions_loan左连接
            merged_temp = pd.merge(
                monthly_balance,
                month_transactions_loan,
                on=['cust_code', 'cust_name', 'spl_code', 'spl_name'],
                how='left'
            ).fillna(0)  # 用0填充缺失值

            # 第二步：将临时结果与month_transactions_repayment左连接
            merged = pd.merge(
                merged_temp,
                month_transactions_repayment,
                on=['cust_code', 'cust_name', 'spl_code', 'spl_name'],
                how='left'
            ).fillna(0)  # 再次用0填充新产生的缺失值

            # 计算回滚余额和借据数（先不处理负数，看看实际情况）
            merged['original_balance'] = merged['current_principal_balance'] + merged['month_repayment_amount'] - \
                                         merged['month_loan_amount']

            # 暂时不处理负数，保留原始值进行观察
            merged['rollback_principal_balance'] = merged['original_balance']  # .clip(lower=0) 暂时注释掉

            # # 检查负余额并标记异常
            # merged['is_negative_balance'] = (merged['rollback_principal_balance'] < 0).astype(int)
            # merged['is_negative_count'] = (merged['rollback_loan_count'] < 0).astype(int)

            # 记录结果（包含异常标记、原始值和集团信息）
            result = merged[['cust_code', 'cust_name', 'spl_code', 'spl_name', 'current_principal_balance',
                             'month_repayment_amount', 'month_loan_amount']].copy()
            result['rollback_date'] = target_date.strftime('%Y-%m-%d')

            # 获取集团映射
            company_mapping_code = DataConfig.get_company_code(conn_dwd)
            print("映射集团客户前数据量：", len(result))
            company_result = add_group_mapping(result, company_mapping_code)
            print("映射集团客户后数据量：", len(company_result))

            # 准备下一个月的数据（使用非负值继续计算）
            monthly_balance = merged[['cust_code', 'cust_name', 'spl_code', 'spl_name', 'rollback_principal_balance']].copy()
            monthly_balance.columns = ['cust_code', 'cust_name', 'spl_code', 'spl_name', 'current_principal_balance']

            results.append(company_result)

            print(f"处理了 {len(result)} 个客户的回滚数据")

    # 合并所有结果
    final_result = pd.concat(results, ignore_index=True)

    # %% 结果处理与输出
    # 排序和格式化
    final_result = final_result.sort_values(['cust_name', 'spl_name', 'rollback_date'])
    final_result['current_principal_balance'] = final_result['current_principal_balance'].round(2)

    # 显示统计信息
    print("\n=== 月倒退回滚结果统计 ===")
    print(f"总记录数: {len(final_result)}")
    print(f"客户数: {final_result['cust_code'].nunique()}")
    print(f"回滚月份: {final_result['rollback_date'].nunique()}")

    print("\n前20条记录:")
    print(final_result.head(20))

    return final_result

    # # 异常数据分析
    # negative_cases = final_result[
    #     (final_result['is_negative_balance'] == 1) |
    #     (final_result['is_negative_count'] == 1)
    # ].copy()
    #
    # print(f"\n=== 异常数据发现 ===")
    # print(f"负余额客户数: {len(negative_cases)}")
    # if len(negative_cases) > 0:
    #     print("负余额详情:")
    #     print(negative_cases[['cust_name', 'rollback_date', 'original_balance', 'original_count',
    #                         'is_negative_balance', 'is_negative_count']])
    #
    # # 按月份汇总（客户维度）
    # monthly_summary = final_result.groupby('rollback_date').agg({
    #     'rollback_principal_balance': 'sum',
    #     'cust_name': 'nunique',
    #     'rollback_loan_count': 'sum',
    #     'is_negative_balance': 'sum',  # 负余额客户数
    #     'is_negative_count': 'sum'     # 负借据数客户数
    # }).reset_index()
    #
    # # 重命名异常列
    # monthly_summary = monthly_summary.rename(columns={
    #     'is_negative_balance': 'negative_balance_customers',
    #     'is_negative_count': 'negative_count_customers'
    # })
    #

    #
    # #%% 与方法1结果对比（可选）
    # print("\n=== 与方法1对比 ===")
    # print("方法2已完成，如需与方法1对比，请运行balance_rollback.py")
    #
    # %%

def calculate_months_to_rollback(current_date: datetime, rollback_year: int) -> int:
    """
    计算从当前日期回滚到目标年份（rollback_year）1月所需的月份数

    Args:
        rollback_year: 目标回滚年份（如2025）
    Returns:
        回滚所需的月份数（如当前是2025年9月，回滚到2025年1月，由于需要最新日期，所以无需-1）
    """
    # 1. 获取当前日期（年、月）
    current_year = current_date.year
    current_month = current_date.month

    # 2. 校验目标年份合理性（避免目标年份晚于当前年份，导致回滚月份为负数）
    if rollback_year > current_year:
        raise ValueError(f"目标年份 {rollback_year} 不能晚于当前年份 {current_year}")
    elif rollback_year == current_year:
        # 若目标年份=当前年份，回滚到当年1月，月份数=当前月份-1
        months_to_rollback = current_month
    else:
        months_to_rollback = (current_year - rollback_year) * 12 + current_month

    return months_to_rollback

def cal_rollback_months(curr_date, roll_year):
    # 计算需要回滚的月份数
    months_to_rollback = calculate_months_to_rollback(curr_date, roll_year)
    print(f"从 {curr_date} 回滚到 {roll_year} 年1月，需回滚 {months_to_rollback} 个月")

    # 动态生成回滚日期列表（包含当前实际日期）
    rollback_months = []

    # 第1个日期：当前实际日期（一般都不是月底）
    rollback_months.append(curr_date.date())

    # 后续日期：每个月的最后一天（倒推）
    for i in range(1, months_to_rollback):
        # 从当前月份倒推i个月
        current_year = curr_date.year
        current_month = curr_date.month

        # 计算目标月份（倒推i个月）
        target_month = current_month - i
        target_year = current_year

        # 处理跨年的情况
        if target_month <= 0:
            target_month += 12
            target_year -= 1

        # 获取该月的最后一天
        if target_month == 12:
            next_month_first = datetime(target_year + 1, 1, 1)
        else:
            next_month_first = datetime(target_year, target_month + 1, 1)

        month_end = next_month_first - timedelta(days=1)
        rollback_months.append(month_end.date())

    # 按时间升序排序（从远到近）
    rollback_months = sorted(rollback_months)

    print(f"当前实时日期: {curr_date.strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"生成的回滚日期: {[d.strftime('%Y-%m-%d') for d in rollback_months]}")
    print(f"最新数据日期: {rollback_months[-1].strftime('%Y-%m-%d')} (当前实际日期)")

    print(f"回滚期间: {rollback_months}")

    return rollback_months


def add_group_mapping(final_result, company_mapping_code):
    # 确保用于关联的字段没有重复值，避免join后数据发散
    # 只保留每个客户最新/唯一的集团映射关系
    unique_mapping = company_mapping_code.drop_duplicates(
        subset=['cust_code', 'cust_name'],  # 按客户唯一标识去重
        keep='last'  # 保留最后一条映射记录
    )

    # 左连接：只保留final_result的所有记录，匹配集团映射
    merged = pd.merge(
        final_result,
        unique_mapping[['cust_code', 'cust_name', 'inv_company_code', 'inv_company_name']],
        on=['cust_code', 'cust_name'],  # 用客户编码和名称双重关联，确保准确性
        how='left'  # 左连接保证final_result的行数不变（不发散）
    )

    # 处理空值：如果集团编码/名称为空，用客户自身的编码/名称填充
    merged['inv_company_code'] = merged['inv_company_code'].fillna(merged['cust_code'])
    merged['inv_company_name'] = merged['inv_company_name'].fillna(merged['cust_name'])

    return merged


def adjust_negative_balances(result_list):
    # 存储所有修改记录的列表
    change_records = []
    final_result = result_list.copy()

    # 遍历每一个DataFrame
    for df_idx, df in enumerate(final_result):
        # 创建副本避免修改原数据
        modified_df = df.copy()

        # 定位负数所在行
        negative_rows = modified_df['current_principal_balance'] < 0

        if negative_rows.any():
            # 提取变化前的数据
            changed_data = modified_df[negative_rows].copy()
            # 记录原始值和来源索引
            changed_data['original_value'] = changed_data['current_principal_balance']
            changed_data['source_df_index'] = df_idx  # 标记来自第几个DataFrame

            # 将负数改为0
            modified_df.loc[negative_rows, 'current_principal_balance'] = 0

            # 更新列表中的DataFrame
            final_result[df_idx] = modified_df

            # 收集变化记录
            change_records.append(changed_data)

    # 合并所有变化记录
    if change_records:
        all_changes = pd.concat(change_records, ignore_index=True)
        # 显示变化的关键信息
        print(f"共修改 {len(all_changes)} 条负余额记录")
        return final_result, all_changes
    else:
        print("未发现需要修改的负余额记录")
        return final_result, None

def write_to_db(engine, final_result, filter_data):
    """
            将final_result中的3个DataFrame写入数据库表
            final_result: 列表，存放3个分类的DataFrame（顺序与filter_data对应）
            filter_data: 列表，数据分类标识（['全量', '优合', '非金融仓']）
            engine: SQLAlchemy创建的数据库连接引擎
            """

    # --------------------------
    # 1. 为每个DataFrame添加“数据分类”列，并统一字段顺序（与表结构对齐）
    # --------------------------
    df_list = []
    # 遍历final_result和filter_data（按顺序对应）
    for df, classify in zip(final_result, filter_data):
        # 复制DataFrame避免修改原数据
        temp_df = df.copy()
        # 添加数据分类列（与表结构中的data_classify字段对应）
        temp_df['data_classify'] = classify
        # 添加创建人（默认Veta，与表结构的insert_name默认值一致，也可显式指定）
        temp_df['insert_name'] = 'Veta'
        # 添加更新时间（当前时间，与表结构的etl_time字段对应）
        temp_df['etl_time'] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')


        # 确保DataFrame的字段顺序与表结构完全一致（避免字段错位）
        table_columns = [
            'data_classify', 'cust_code', 'cust_name', 'spl_code', 'spl_name', 'rollback_date',
            'current_principal_balance', 'month_repayment_amount', 'month_loan_amount',
            'inv_company_code', 'inv_company_name', 'insert_name', 'etl_time'
        ]
        # 筛选并排序字段（若DataFrame缺少某字段，会填充为NaN，需确保final_result包含核心字段）
        temp_df = temp_df[table_columns]

        df_list.append(temp_df)

    # --------------------------
    # 2. 拼接所有DataFrame（纵向拼接，无数据发散）
    # --------------------------
    combined_df = pd.concat(df_list, ignore_index=True)
    # 可选：处理空值（如集团编码/名称为空时用客户自身编码/名称填充，与之前逻辑一致）
    # combined_df['inv_company_code'] = combined_df['inv_company_code'].fillna(combined_df['cust_code'])
    # combined_df['inv_company_name'] = combined_df['inv_company_name'].fillna(combined_df['cust_name'])

    # --------------------------
    # 3. 清空表中所有数据（先删后插，避免重复）
    # --------------------------
    with engine.connect() as conn:
        # 执行TRUNCATE（清空表，重置自增ID；若表无自增，用DELETE FROM也可）
        # 注意：TRUNCATE无法回滚，若需事务安全，可改用DELETE FROM
        truncate_sql = text("TRUNCATE TABLE dwh_tmp.tmp_customer_rollback_balance_special;")
        conn.execute(truncate_sql)
        # SQLAlchemy自动处理事务提交，无需手动commit()
        # conn.commit()

    # --------------------------
    # 4. 将拼接后的数据写入数据库表
    # --------------------------
    combined_df.to_sql(
        name='tmp_customer_rollback_balance_special',  # 表名
        schema='dwh_tmp',  # 数据库 schema（即库名）
        con=engine,  # 数据库连接
        if_exists='append',  # 追加模式（因已清空表，append等同于覆盖）
        index=False,  # 不写入DataFrame的索引（表中无索引字段）
        chunksize=1000  # 批量插入（数据量大时避免内存溢出）
    )

    print(f"数据写入完成！共写入 {len(combined_df)} 条记录。")
    return combined_df

def export_excel(filter_data, final_result):

    # 按集团和月份汇总
    company_balance_result = []
    for i in range(len(final_result)):
        company_summary = final_result[i].groupby(['rollback_date', 'inv_company_code', 'inv_company_name', 'spl_code', 'spl_name']).agg({
            'current_principal_balance': 'sum'
        }).reset_index()
        company_balance_result.append(company_summary)

    if not company_balance_result:
        print("集团客户维度聚合成功！")

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_file = f'2025年南美垫资回滚结果_月倒退法_{timestamp}.xlsx'
    #
    with pd.ExcelWriter(output_file, engine='openpyxl') as writer:
        for i in range(len(final_result)):
            # 1. 处理客户回滚明细表：替换列名为中文
            df_customer = final_result[i].copy()
            # 只保留映射中存在的列（避免多余列），并替换为中文
            df_customer = df_customer.rename(columns=DataConfig.FIELD_MAPPING)
            # 导出客户明细表
            df_customer.to_excel(
                writer,
                sheet_name=f'{filter_data[i]}数据客户回滚明细',
                index=False
            )

            # 2. 处理集团月度汇总表：替换列名为中文
            df_company = company_balance_result[i].copy()
            df_company = df_company.rename(columns=DataConfig.FIELD_MAPPING)
            # 导出集团汇总表
            df_company.to_excel(
                writer,
                sheet_name=f'{filter_data[i]}数据集团客户月度汇总',
                index=False
            )
            # 月度汇总
            # monthly_summary.to_excel(writer, sheet_name='月度汇总', index=False)
            # 异常数据
            # if len(negative_cases) > 0:
            #     negative_cases.to_excel(writer, sheet_name='异常数据明细', index=False)

    print(f"\n结果已保存到: {output_file}")

def main(environment, db_type, db_name):
    print('\nStart\n')

    # 数据库连接
    print('获取数据库连接\n')

    # environment = 'local'
    conn_dwd = db_connector.conn_db(environment, db_type, db_name)

    print('获取回滚月份末日期\n')

    # 获取当前日期和回滚期间日期
    current_date = datetime.now()  # 实时获取当前日期
    rollback_year = 2025

    rollback_months = cal_rollback_months(current_date, rollback_year)


    # 获取各分类客户垫资回滚数据
    filter_data = ['全量', '优合', '非金融仓']
    # filter_data = ['全量']
    finance_type = [0]    # 0-不分融资类型；1-分融资类型

    final_result = []
    for f in finance_type:
        for l in filter_data:

            query = DataQuery(conn_dwd)
            current_balance = query.query_current_balance(f, l)

            # 获取当前客户本金余额
            print(f"当前时点客户数: {len(current_balance)}")
            print("当前余额示例:")

            print(f"\n开始 {l} 数据月倒退回滚计算...\n")

            result = calculate_rollback(current_balance, current_date, rollback_months, conn_dwd, f, l)

            print(f"\n {l} 数据月倒退回滚计算完成！\n")

            final_result.append(result)

    # 数据清洗，负数填充0
    print('检查本金余额为负的情况：')
    pre_final_result, changes_list = adjust_negative_balances(final_result)
    print(changes_list)

    # 导入数据库
    print('\n写入数据库')
    write_to_db(conn_dwd, pre_final_result, filter_data)
    print('\n写入成功！')

    # 导出excel
    # print("\n按集团客户维度聚合数据并导出Excel\n")
    # export_excel(filter_data, pre_final_result)





if __name__ == '__main__':
    main(environment='local', db_type='production', db_name='dwh_tmp')
